import keras
from keras import layers
# 图像数据集预处理
from keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt

# 检索训练集 train 以及验证集 validation
train_dir = "./data/numbers/train"
validation_dir = "./data/numbers/validation"

# 对图像进行预处理
train_datagen = ImageDataGenerator(
    rescale=1./255 # 将所有的像素值控制到 0 - 1之间（减少后期训练计算量）
)
validation_datagen = ImageDataGenerator(
    rescale=1./255
)

# 构建图像的数据集
train_generator = train_datagen.flow_from_directory(
    train_dir,
    target_size=(224,224),
    class_mode="categorical" # 文件夹的方式进行分类，目前是0，1，2三分类
)
validation_generator = validation_datagen.flow_from_directory(
    validation_dir,
    target_size=(224,224),
    class_mode="categorical"
)

# ------------------构建一个简易的卷积神经网络------------------

model = keras.Sequential([
    # 图像输入
    layers.Conv2D(32,(3,3),activation="relu",input_shape=(224,224,3)),
    layers.MaxPool2D(pool_size=(2,2),strides=2),
    # 隐藏层
    layers.Conv2D(64,(3,3),activation="relu",strides=1),
    layers.MaxPool2D(pool_size=(2,2),strides=2),
    # 图像拉伸成平直的一维数据
    layers.Flatten(),
    # 全连接 + 输出
    layers.Dense(1000),
    layers.Dense(3,activation="softmax")
])

model.compile(
    optimizer="adam",
    loss=keras.losses.binary_crossentropy,
    metrics=['acc']
)

history = model.fit(train_generator,
                    epochs=20,
                    validation_data=validation_generator,
                    shuffle=True,
                    verbose=1)

model.save('temp.keras')

loss = history.history["loss"]
acc = history.history["acc"]
val_loss = history.history["val_loss"]
val_acc = history.history["val_acc"]

epochs = range(len(acc))

plt.plot(epochs,acc,'r',label="Traing accuracy")
plt.plot(epochs,val_acc,'b',label="Validation accuracy")
plt.plot(epochs,loss,'r--',label="Traing loss")
plt.plot(epochs,val_loss,'b--',label="Validation loss")
plt.title('Training and Validation accuracy')
plt.legend(loc=0)
# plt.figure()

plt.show()














